from operator import mod
from pyexpat import model
import torchvision.models as models
import torch
from torchsummary import summary

def count_param(model):
    param_count = 0
    for param in model.parameters():
        param_count += param.view(-1).size()[0]
    return param_count

if __name__ == '__main__':
    resnet18 = models.resnet18()
    resnet18_size = count_param(resnet18)
    resnet18_msg = '{:.2f}M ({})'.format(
        resnet18_size/1024**2, resnet18_size)
    print(resnet18_msg)

    num_parameters = sum(torch.numel(parameter) for parameter in resnet18.parameters())
    print(num_parameters)

    summary(resnet18.cuda(), (3,64,64))